import torch
import torch.utils.data as Data
from dataset import TestDataset, BookUserInfo
from tqdm import tqdm
import json
from typing import Dict, List
from pathlib import Path


class Evaluator:
    def __init__(self, hyper_params, book_user_info: BookUserInfo):
        self.hyper_params = hyper_params
        
        self.dataset = TestDataset(hyper_params, book_user_info)
        self.label_warm = json.load(open(Path(hyper_params['data_path'])/'interaction_data/interaction_test.json', 'rt'))
        self.label_cold = json.load(open(Path(hyper_params['data_path'])/'interaction_data/interaction_test_cold.json', 'rt'))

    def __evaluate(self, result: Dict[int, List[int]], recall_rank=[1, 5, 20, 50, 100, -1], cold_start=False):
        label = self.label_cold if cold_start else self.label_warm

        total_cnt = 0
        total_hit = [0] * len(recall_rank)

        for user_id, interact_seq in label.items():
            user_id = int(user_id)

            for interact in interact_seq:
                total_cnt += 1
                book_id = interact['book_id']

                for i, rank in enumerate(recall_rank):
                    if rank > 0 and book_id in result[user_id][:rank]:
                        total_hit[i] += 1
                    elif rank <= 0 and book_id in result[user_id]:
                        total_hit[i] += 1

        for i in range(len(recall_rank)):
            print(f'Hit@{recall_rank[i]}: {total_hit[i]}/{total_cnt} HR@{recall_rank[i]}: {total_hit[i]/total_cnt}')


    def evaluate(self, model):
        print('evaluate: Evaluating...')
        result = [None] * self.hyper_params['user_cnt']

        loader = Data.DataLoader(self.dataset, self.hyper_params['eval_batch_size'], num_workers=16)
        model.eval()

        result_tensor = []
        result_bookid = []

        with torch.no_grad():
            for data in loader:
                user_id, school, grade, pos_bid, pos_title, pos_type = [i.to(self.hyper_params['device']) for i in data]

                score = model(user_id, school, grade, pos_bid, pos_title, pos_type)
                result_tensor.append(score.detach().cpu())
                result_bookid.append(pos_bid.detach().cpu())

        result_tensor = torch.cat(result_tensor, dim=0).reshape(self.hyper_params['user_cnt'], -1)
        result_bookid = torch.cat(result_bookid, dim=0).reshape(self.hyper_params['user_cnt'], -1)

        result_tensor = torch.argsort(result_tensor, dim=-1, descending=True)
        result_bookid = torch.gather(result_bookid, dim=1, index=result_tensor)

        # Convert to result
        for i in range(self.hyper_params['user_cnt']):
            result[i] = result_bookid[i].tolist()

        print('Warm start result: ')
        self.__evaluate(result, cold_start=False)
        print('Cold start result:')
        self.__evaluate(result, cold_start=True)

        model.train()
